'''
Michael Cabot

To run this program you need to install the package OpenCV at:
http://opencv.willowgarage.com/wiki/InstallGuide
The documentation can be found at:
http://opencv.willowgarage.com/documentation/python/index.html


PSEUDO CODE:
    Resize the taken image to (160, 120) and    >> convert()
        convert the type to a cv.IplImage
    Mark all the blue/yellow pixels             >> markGoalCV()
    Blur the image                              >> blur()
    Create 2 groups containing the most left    >> findGoal()
        and most right pixels
    Calculate the threshold to remove outliers  >> threshold()
    If 1 of the 2 groups has less than 5        >> EXIT
        elements, then return 'None'
    Remove the outliers from the 2 groups       >> removeOutliers
    Find the 2 tops of the 2 groups             >> findTop()
    Calculate the width of the 2 goal poles     >> findWidth()
    Calculate the distance and angles to the    >> distance()
        2 poles
    Calculate the position of the 2 poles       >> calcPosition
    If the distance to the position between     >> EXIT
        the 2 goals is greater than 6 meters,
        than return 'None'
    Devide the x- and y-coordinate of the       >> EXIT
        centre of the goal by 1000 (convert
        from milimeters to meters) and switch
        them (x <--> y).
        Return the position and angle of the
        goal.


                                                                                                  ^_^
UPDATED (28 march by Anna))
added markGoalCV(im, color) which filters the image with openCV instead of PIL (replaces markGoal, checkGoal)
we don't do anything with PIL now so the image it expects is a resized openCV image

TODO (28 march by Anna)
- Make a HSV filter for a blue goal
- Get the width and height with openCV instead of PIL | done
- Fix snipePosition (doesn't work with all the cv now)

Bij findWidth loopen we nu door het plaatje heen op een PIL manier, maar dat kan nu niet meer,
want het plaatje is met openCV gefilterd.

UPDATED (1 april by Anna):
Discovered that removing the outliers did not work really well. Problem: we removed the outliers while still searching for
outliers which messed with the size of the group.

UPDATE (2 april by Michael):
-I've updated the comments above.
-removed the line:
        meanLeft = findMean(leftgroup)
which wasn't being used.
- now use distance() in stead of calcWidth() which did the same thing.
   Note that the radius R is changes accordingly to the radius of the pole
- change the last few steps in the code: calclulate the position of the
    2 poles. Then the position of the goal is its average.

UPDATE (3-04 by Michael)
Removed the imports of Imagefilter, ImageEnhance and ImageOps since
they weren't been used.
Because the houghCircles ball tracker might need a different size image
than the goal tracker the snapshot function only passes the taken image.
LocateGoal_04_CV has to reshape (160 by 120) and convert (to OpenCV) the image.
The function snipePosition() has been put in comments (see OLD STUFF below)

UPDATE (morning 4 april by Anna)
-Discovered that the distance calculation is not correct.(it thinks a goal is at 6m when it should be +-2m.
-The filter for the blue goal wasn't implemented before and getting the HSV values right is not going well so far.
-We've looked up the rgb colours of the official goal which are now in rgb&hsv_goal

UPDATE (afternoon 4 april by Anna)
Trying to make the code faster.
-made 1 for loop instead of 2 separate in calcWidth() and also did that in findGoal()
-removed all the images that were saved
-removed the prints
-removed all the position calculations since we don't use them now (they are incorrect)

UPDATE (5 april (set up Iran Open) by Anna)
-made pictures of the yellow goal first and when trying to make the filter discovered that the HSV values go from:
H -> 0-180
S -> 0-255
V -> 0-255
instead of 0-255 as we thought initially.
The yellow goalfinder now works almost perfect except at the other side of the field.
-also made pictures of the blue goal, but the blue is almost the same as the black under the boarding, so that will be more difficult.
The hsv values are in the dropbox.
-it turned out I couldn't combine two loops into one (they found the groups)
-the blue goal in Iran is too dark and will be painted lighted, so we have to adjust our blue filter tomorrow.
'''

from PIL import Image
import math
import cv


width = None                                    #the width of the picture (is set in 'markOrange')
height = None                                   #the height of the picture (is set in 'markOrange')


#snapshots on which I tested this code
#photos = ["snapShot1", "snapShot2", "snapShot3", "snapShot4", "snapShot11", "snapShot12", "snapShot13", "snapShot14", "snapShot15", "snapShot20"]

photo = None                          #photo taken by the Nao. Is set when in filter() or when consulting this program.
size = None

def run(im, color, head):
    im = convertImage(im)
    [yawHead, pitchHead] = head
    #cv.SaveImage('imageBefore'+str(color)+'.jpg', im)
    
    #print 'Mark Orange...'
    im = markGoalCV(im, color)  #filter image to find all the blue/yellow pixels (makes them white)
    #cv.SaveImage('filterImage_CV'+str(color)+'.jpg', im)    
    #print 'Blur..'
    #print 'before size ' + str(width) + ' ' + str(height)
    blurredImage = blur(im)         # blur the image
    #cv.SaveImage('blurImage_CV'+str(color)+'.jpg', blurredImage)
    
    (leftgroup, rightgroup) = findGoal(blurredImage)
    
    #print 'leftgroup with outliers'
    #print leftgroup
    #print 'rightgroup with outliers'
    #print rightgroup

    # return None if the left or right groups has less than 5 elements
    if len(leftgroup) < 5 or len(rightgroup) < 5:
        return None

    threshold = findThreshold( leftgroup, rightgroup)
    leftgroup = removeOutliers(leftgroup, threshold)
    rightgroup = removeOutliers(rightgroup, threshold)

    #print 'left without outliers'
    #print leftgroup
    #print 'right without outliers'
    #print rightgroup
    
    #print "threshold " + str(threshold)
    leftTop = findTop(leftgroup, threshold)
    rightTop = findTop(rightgroup, threshold)
    #print "leftTop " + str(leftTop)
    #print "rightTop " + str(rightTop)
    
    (leftX, leftY) = leftTop
    (rightX, rightY) = rightTop
    mid = ((leftX+rightX)/2, (leftY+rightY)/2)
    (xcoord, ycoord) = mid
    angleToGoal = calcXAngle(xcoord, yawHead)
    #print 'angle to goal ' + str(color) + ': ' + str(angleToGoal)
    print 'in locateGoal_IR:'
    print angleToGoal
    return angleToGoal

def calcXAngle(xcoord, yawHead):
        (width, height) = size
        xDiff = width/2 - xcoord
        # distanceToFrame = 0.5 * width/math.tan(0.5*horizontalAngle)
        # horizontalAngle = 46.4 degrees = 0.8098327729253689 radians
        # math.tan(0.5*horizontalAngle) = 0.42860054745600146
        distanceToFrame = 0.5 * width / 0.42860054745600146
        xAngle = math.atan(xDiff/distanceToFrame)
        return xAngle + yawHead
    
def convertImage(picture):
    global size
    size = (160, 120)
    # resize the image
    picture = picture.resize(size)
    # convert the type to OpenCV
    image = cv.CreateImageHeader(size, cv.IPL_DEPTH_8U, 3)
    cv.SetData(image, picture.tostring(), picture.size[0]*3)
    return image


def markGoalCV(im, color):
    # Size of the images
    size = (160, 120)
    
    hsvFrame = cv.CreateImage(size, cv.IPL_DEPTH_8U, 3)
    filter = cv.CreateImage(size, cv.IPL_DEPTH_8U, 1)
    
    # Goalfilter values
    if(color == 'blue'): #KLOPT NIET.. ZOEK DE JUISTE HUE VALUES        
        hueMin = 110
        hueMax = 142
        saturationMin = 45
        saturationMax = 222
        valueMin = 110
        valueMax = 240
    if(color == 'yellow'): #NEEDS MORE TESTING
        hueMin = 24
        hueMax = 80
        saturationMin = 45
        saturationMax = 256
        valueMin = 80
        valueMax = 256
    
    hsvMin1 = cv.Scalar(hueMin, saturationMin, valueMin, 0)
    hsvMax1 = cv.Scalar(hueMax, saturationMax, valueMax, 0)
    # Color detection using HSV
    #print 'CvtColor'
    #print im, hsvFrame
    cv.CvtColor(im, hsvFrame, cv.CV_BGR2HSV)
    cv.InRangeS(hsvFrame, hsvMin1, hsvMax1, filter)
    #print 'hsv ' + str(hsvFrame)
    #print 'filter' + str(filter)
    return filter

def blur(cv_im):
    # blur the OpenCV image
    cv.Smooth(cv_im, cv_im, cv.CV_BLUR, 10, 10)
    return cv_im

def findGoal(im):
    global width
    global height
    (width, height) = cv.GetSize(im) #SIZE VIA OPENCV
    leftgroup = []
    rightgroup = []
    #print im[0,0]
    for y in range(height):
        for x in range(width):
            if( im[y,x]/256 > 0.5 ):
                leftgroup.append((x,y))
                #print x, y, im[y,x]
                break
    
    for y in range(height):
        for x in range(width):
            if(im[y, width-x-1]/256 > 0.5 ):
                rightgroup.append((width-x-1, y))
                break
               
    return (leftgroup, rightgroup)

def findTop(group, threshold):
    mean = findMean(group)
    maxY = height
    for i in group:
        (x, y) = i
        if( y < maxY ):
            maxY = y
    return (mean, maxY)


	
def findWidth(im, mid, leftGroup, rightGroup):
    #print "mid" + str(mid)
    (x, y) = mid
    maxHeight = 0
    #finds the bottom of the bar
    
    # als je een witte pixel hebt gevonden en
    # daarna een zwarte pixel vind>> dan heb
    # je de onderkant van de lat gevonden
    #NOTE: moet nog aangepast worden voor als de lat schuin is.
    foundWhite = False
    for i in range(height):
        if not(im[i,x] == 0):
            foundWhite = True
            maxHeight = i
        if im[i,x]==0 and foundWhite:
            break
        
    
    #print "maxHeight " + str(maxHeight)
 
    if not foundWhite:
        print 'could not find the upper bar'
        return None			#if you haven't found a maxHeight, stop.
    
    h = maxHeight + 5 #TRESHOLD to make sure you're below the bar
    #print 'h: ' + str(h)
    (xLeft, yLeft) = leftGroup.pop()
    (xRight, yRight) = rightGroup.pop()
    
    #find the most left point of the left group and the most right one
    leftLeft = None
    leftRight = None
    rightLeft = None
    rightRight = None
    for i in range(1, x):
        if( (xLeft-i > 0) and not(im[h, xLeft-i] == 0) ):
            leftLeft = xLeft-i
        if( (xRight-i > 0) and not(im[h, xRight-i] == 0) ):
            leftRight = xRight-i
        if( (xLeft+i < height) and not(im[h, xLeft+i] == 0) ):
            rightLeft = xLeft+i
        if( (xRight+i < width) and not(im[h, xRight+i] == 0) ):
            rightRight = xRight+i
    #print "leftLeft " + str(leftLeft)
    #print "leftRight " + str(leftRight)

    #print "rightLeft " + str(rightLeft)
    #print "rightRight " + str(rightRight)
    #average width
    #TO DO
    #compare if the width in the image is +- the same
    return (leftLeft, rightLeft, leftRight, rightRight, maxHeight)
	
def findMean(group):
    sumX = 0
    for i in group:
        (x, y) = i
        sumX = sumX + x
    mean = sumX / len(group)
    return mean

def findThreshold(leftGroup, rightGroup):
    (leftMin, leftMax) = findMinMax(leftGroup)
    (rightMin, rightMax) = findMinMax(rightGroup)
    maxX = max(leftMax, rightMax)
    minX = min(leftMin, rightMin)
    #print 'max: ' + str(maxX)
    #print 'min: ' + str(minX)
    dist = maxX - minX
    if( dist < 15 ):
        return 15
    return dist/2 

def findMinMax(group):
    (x,y) = group[0]
    minX = maxX = x
    for i in group:
        (x, y) = i
        if( x > maxX):
            maxX = x
        if( x < minX):
            minX = x
    return(minX, maxX)

def removeOutliers(group, treshold):
    mean = findMean(group)
    outliers = []
    #print 'mean in group ' + str(mean)
    for i in group:
        (x,y) = i
        dist = math.fabs(mean - x)
        if(dist > treshold):
            #print 'outlier' + str(i)
            outliers.append(i)
    #print outliers
    for i in outliers:
        group.remove(i)
    return group

def distance(coord, r):
    (xcoord, ycoord) = coord
    #mmperpixel = R/r
    horizontalangle = 46.4
    verticalangle = 34.8
    global distancetoframe
    R = 90 #radius of the paal in mm's
    distancetoframe = 0.5*width/math.tan(0.5*horizontalangle*math.pi/180)
    xdiff = width/2 - xcoord
    ydiff = height/2 - ycoord
    lx = math.sqrt(distancetoframe**2 + xdiff**2)
    ly = math.sqrt(distancetoframe**2 + xdiff**2)
    l = math.sqrt(lx**2 + ly**2)
    angletoballx = math.atan(xdiff/distancetoframe)*180/math.pi
    angletobally = math.atan(ydiff/distancetoframe)*180/math.pi
    L = R*l/r
    return (L, angletoballx, angletobally)


def calcPosition((L, angletoballx, angletobally), HeadPitch, HeadYaw):
    #distance to camera = sqrt( zOffsetCam^2 + yOffsetCam^2 ) ~= 86.69
    #z Offset to Camera = 67.90
    #y Offset to Camera = 53.90
    distToCam = 86.69

    
    #pitch of the camera = pitchHead + atan(yOffsetCam / zOffsetCam) ~= pitchHead + 0.671
    camPitch = HeadPitch + 0.671
    #print 'camPitch in degrees: ' + str(camPitch* 180 / math.pi)
    #y cam (distance to camera over y-axis) = sin(camPitch) * distToCam
    yCam = math.sin(camPitch) * distToCam
    #print 'yCam: ' + str(yCam)
    #yCamToBall: is the distance over the y-axis from the camera to the ball
    #yCamToBall = sin(theta) * L    where theta is the angle opposite of the camera-position
    #                               and the ball-position over the y-axis
    #theta = pi - camPitch - "opposite angle to z-offset to camera"
    #where "opposite angle to z-offset to camera" = atan(zOffsetCam / yOffsetCam ~= 0.8998
    theta = math.pi - camPitch - 0.8998
    

    #total y-angle is the angle opposite to position of the camera and the ball on the y-axis
    totalYAngle = theta + angletobally
    #print 'totalYAngle in degrees: ' + str(totalYAngle*180 / math.pi)

    #yCamToBall is the distance over the y-axis from the camera to the ball
    yCamToBall = math.sin(totalYAngle) * L
    #print 'yCamToBall: ' + str(yCamToBall)
    #y is the sum of the distance to the camera from the neck and the distance from the camera to the ball
    y = yCam + yCamToBall
    x = math.tan(angletoballx) * y

    # this (x,y) position must be corrected for its HeadYaw (looking to the left or right)
    # this can be done using a rotation matrix:
    # [cos(a), -sin(a); sin(a), cos(a)]
    x = math.cos(HeadYaw)*x + -math.sin(HeadYaw)*y
    y = math.sin(HeadYaw)*x + math.cos(HeadYaw)*y
    
    return (x,y)

'''
def test():
    #im = cv.LoadImage("yellowR.jpg")
    im = Image.open("snapShotYellowGoal.jpg")
    #im = Image.open("blueGoal02.jpg")
    color = 'yellow'
    run(im, color, (0,0))
    #print 'finished'
'''    
#test()
# old stuff
'''
def calcWidth(coord, r):
    (xcoord, ycoord) = coord
    #mmperpixel = R/r

    horizontalangle = 46.4
    verticalangle = 34.8
    global distancetoframe
    R = 90 #radius of the paal in mm's
    distancetoframe = 0.5*width/math.tan(0.5*horizontalangle*math.pi/180)
    xdiff = width/2 - xcoord
    ydiff = height/2 - ycoord
    lx = math.sqrt(distancetoframe**2 + xdiff**2)
    ly = math.sqrt(distancetoframe**2 + xdiff**2)
    l = math.sqrt(lx**2 + ly**2)
    angletoballx = math.atan(xdiff/distancetoframe)*180/math.pi
    angletobally = math.atan(ydiff/distancetoframe)*180/math.pi
    L = R*l/r
    return (L, angletoballx, angletobally)


    #to test if the centre of the ball is correct
def snipePosition((x, y)):
    global pix
    for i in range(width):
        pix[i, y] = (0, 255, 0)

    for i in range(height):


        pix[x, i] = (0, 255, 0)
    return pix


##############!!!!!!!!!!!!!!!!!##############
# de oude 'run()' waarbij de position wordt berekend

def run(im, color, head):
    im = convertImage(im)
    [yawHead, pitchHead] = head
    cv.SaveImage('imageBefore'+str(color)+'.jpg', im)
    
    #print 'Mark Orange...'
    im = markGoalCV(im, color)  #filter image to find all the blue/yellow pixels (makes them white)
    cv.SaveImage('filterImage_CV'+str(color)+'.jpg', im)    
    #print 'Blur..'
    #print 'before size ' + str(width) + ' ' + str(height)
    blurredImage = blur(im)         # blur the image
    cv.SaveImage('blurImage_CV'+str(color)+'.jpg', blurredImage)
    
    (leftgroup, rightgroup) = findGoal(blurredImage)
    
    #print 'leftgroup with outliers'
    #print leftgroup
    #print 'rightgroup with outliers'
    #print rightgroup

    # return None if the left or right groups has less than 5 elements
    if len(leftgroup) < 5 or len(rightgroup) < 5:
        return None

    threshold = findThreshold( leftgroup, rightgroup)
    leftgroup = removeOutliers(leftgroup, threshold)
    rightgroup = removeOutliers(rightgroup, threshold)

    #print 'left without outliers'
    #print leftgroup
    #print 'right without outliers'
    #print rightgroup
    
    #print "threshold " + str(threshold)
    leftTop = findTop(leftgroup, threshold)
    rightTop = findTop(rightgroup, threshold)
    #print "leftTop " + str(leftTop)
    #print "rightTop " + str(rightTop)
    
    (leftX, leftY) = leftTop
    (rightX, rightY) = rightTop
    mid = ((leftX+rightX)/2, (leftY+rightY)/2)
    # find the left and right border of the left and right pole
    width = findWidth(blurredImage, mid, leftgroup, rightgroup)
    
    if width == None:
        print "Could not find 1 of the result of 'findWidth'"
        print "Check the filtered image."
        return None
    
    (lL, rL, lR, rR, maxHeight) = width
	
    # check if 1 of the arguments are None
    if lL==None or rL==None or rR==None or maxHeight==None:
        print "Could not find 1 of the result of 'findWidth'"
        print "Check the filtered image."
        return None
    
	
	
    widthLeft = rL - lL
    widthRight = rR - lR
    #print "width bars:"
    #print "widthLeft: " + str(widthLeft)
    #print "widthRight: " + str(widthRight)

    distanceLeft = distance(leftTop, widthLeft)
    distanceRight = distance(rightTop, widthRight)
    #print "distanceLeft: " + str(distanceLeft)
    #print "distanceRight: " + str(distanceRight)
    ### no longer used (not correct)
    # calculate the position of the left and right pole
    #print 'Calculate Position...'
    #(xPosLeft, yPosLeft) = calcPosition(distanceLeft, pitchHead, yawHead)
    #(xPosRight, yPosRight) = calcPosition(distanceRight, pitchHead, yawHead)
    #print (xPosLeft, yPosLeft)
    #print (xPosRight, yPosRight)
    
    #xPosition = (xPosLeft + xPosRight) / 2
    #yPosition = (yPosLeft + yPosRight) / 2
    #position = (xPosition, yPosition)
    #L = math.sqrt(xPosition**2 + yPosition**2)
    # if the distance to the goal is greater than 6 meters
    #   than it isn't the goal and return None
    #if (L > 6000):
    #    return None
    # the motion module wants a coordinate (x,y) in meters
    # the following converts the position from milimeters to meters (divide by 1000)
    #(x,y) = position
    #x = x/1000
    #y = y/1000
    # the motion module uses x as front-back and y as left-right
    #position = (y, x)

    (LL, xAngleL, yAngleL) = distanceLeft
    (LR, xAngleR, yAngleR) = distanceRight
    xAngle = (xAngleL + xAngleR) / 2
    #yAngle = (yAngleL + yAngleR) / 2
    #angle = (xAngle, yAngle)
    
    #print position
    #print xAngle
    #print yAngle
    return xAngle+yawHead

'''
